import json
import os
import random
import shutil
import math

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
import torch.utils.data.distributed
import torch.backends.cudnn
import torchvision.transforms as transforms
from timm.loss import SoftTargetCrossEntropy
import timm
import cv2
import h5py
from models.BiConv.BiConv import *

import warnings
warnings.filterwarnings("ignore")
matplotlib.use('Agg')
torch.backends.cudnn.benchmark = False

DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
weight = r'/home/yeadc/Documents/cxj/Classify-Seg/task_classify/task/checkpoints/biconv_base_240423_x1q1k2v2/model_140_90.365.pth'

model = torch.load(weight)
model.eval()
model.to(DEVICE)
print(model)

h5file = r'/home/yeadc/Documents/cxj/Classify-Seg/task_classify/task/heatmap/Bump_FY-C201F_MD_ASI_lFFK100951_w1_880_Topography3.h5'
f = h5py.File(h5file, 'r')  # 打开h5文件
class_tiff = f['class_tiff'][:]
defect_tiff = f['defect_tiff'][:]
f.close()

transform_test = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.40272543], std=[0.13901867])
])

image_array = np.array(class_tiff) / 255.
image_class = Image.fromarray(class_tiff).convert("L")
image_class = transform_test(image_class)
image_class = torch.unsqueeze(image_class, dim=0).to(DEVICE)

image_defect = Image.fromarray(defect_tiff).convert("L")
image_defect = transform_test(image_defect)
image_defect = torch.unsqueeze(image_defect, dim=0).to(DEVICE)

# plt.figure(figsize=(8, 8))
# plt.imshow(image_array, cmap='gray')
# plt.axis('off')
# plt.show()



activation = {}
def get_activation(name):
    def hook(model, input1, input2, output):
        activation[name] = output.detach()
    return hook


model.cbam4_lf.register_forward_hook(get_activation('cbam4_lf'))
_ = model(image_class, image_defect)
cbam4_lf = activation['cbam4_lf']
cbam4_lf = cbam4_lf.cpu().numpy()


# plt.figure(figsize=(8, 8))
# plt.imshow(cbam4_lf[0][0], cmap='gray')
# plt.axis('off')
# plt.show()



